from avalanche.evaluation import Metric
import torch
from torch import Tensor
from typing import List, Union


class cusMetric(Metric):
    """Custom Metric

    Args:
        Metric (Metric): Metric
    """

    def __init__(self, name: str = "") -> None:
        self.name: str = name
        self.summed: float = 0.0
        self.weight: float = 0.0

    def result(self):
        if self.weight == 0.0:
            return 0.0
        return self.summed / self.weight

    def reset(self):
        """
        Resets the metric.

        :return: None.
        """
        self.summed = 0.0
        self.weight = 0.0

    def update(self, value, weight=1.0):
        value = float(value)
        weight = float(weight)
        self.summed += value * weight
        self.weight += weight

    def __add__(self, other: "cusMetric") -> "cusMetric":
        """
        Return a metric representing the weighted mean of the 2 means.

        :param other: the other mean
        :return: The weighted mean"""
        res = cusMetric()
        res.summed = self.summed + other.summed
        res.weight = self.weight + other.weight
        return res

    @property
    def value(self) -> float:
        if self.weight == 0.0:
            return 0.0
        return self.summed / self.weight

    @property
    def strs(self) -> str:
        s = f"{self.value:.3e}"
        return s


class accMetric(cusMetric):
    def __init__(self, name: str = "Accuracy") -> None:
        super().__init__(name)

    @torch.no_grad()
    def p_update(self, value, weight=1.0):
        super().update(value, weight)

    @torch.no_grad()
    def update(
        self,
        predicted_y: Tensor,
        true_y: Tensor,
    ) -> None:
        """Update the running accuracy given the true and predicted labels.

        :param predicted_y: The model prediction. Both labels and logit vectors
            are supported.
        :param true_y: The ground truth. Both labels and one-hot vectors
            are supported.

        :return: None.
        """
        true_y = torch.as_tensor(true_y)
        predicted_y = torch.as_tensor(predicted_y)

        if len(true_y) != len(predicted_y):
            raise ValueError(
                "Size mismatch for true_y and predicted_y tensors"
            )

        # Check if logits or labels
        if len(predicted_y.shape) > 1:
            # Logits -> transform to labels
            predicted_y = torch.max(predicted_y, 1)[1]

        if len(true_y.shape) > 1:
            # Logits -> transform to labels
            true_y = torch.max(true_y, 1)[1]

        true_positives = float(
            torch.sum(torch.eq(predicted_y, true_y))
        )
        total_patterns = len(true_y)
        super().update(
            true_positives / total_patterns, total_patterns
        )

    @property
    def strs(self) -> str:
        s = f"{self.value*100:.2f}"
        return s


class Metrics:
    """Format multi Metrics

    Attributes:
        metrics (List[cusMetric]): list of metrics
        format_str (str): format string
    """

    def __init__(self, metrics: List[cusMetric], prefix="  ") -> None:
        self.prefix = prefix
        self.metrics = metrics

    def update(
        self, values: List, weights: Union[List[float], int] = 1
    ) -> None:
        weights_list = (
            [weights] * len(values)
            if isinstance(weights, int)
            else weights
        )
        for m, v, w in zip(self.metrics, values, weights_list):
            m.update(v, w)

    def reset(self) -> None:
        for m in self.metrics:
            m.reset()

    def display(self) -> str:
        print(self.format_str)
        return self.format_str

    @property
    def format_str(self) -> str:
        if len(self.metrics) > 1:
            strs = " ".join(
                [
                    self._vivid(self.prefix) + "\n",
                    *[
                        (" " * 8 + f"{i.name:<20s}: {i.strs}\n")
                        for i in self.metrics
                    ],
                ]
            )
        else:  # only have one metric
            strs = self._vivid(
                self.prefix
                + "   "
                + f"{self.metrics[0].name:<20s}: {self.metrics[0].strs}"
            )

        return strs

    def __getitem__(self, idx: int) -> cusMetric:
        return self.metrics[idx]

    def title(self, title_str) -> str:
        return self._vivid(" " * 10 + title_str + " " * 10)

    @classmethod
    def _vivid(cls, strs):
        border = "-" * len(strs)
        formatted_string = f"\n{border}\n{strs}\n{border}"
        return formatted_string
